package cn.com.main;

import cn.com.consts.MsgCd;
import cn.com.utils.WebUtil;
import cn.hutool.core.io.IoUtil;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Base64;

/**
 * rsa加密示例
 *
 * @author LiuGuodong
 * @date 2024/03/06
 */
@Slf4j
public class RSAEncryptionExample {

    public static final String RSA = "RSA";

    public static void main(String[] args) {
        try {
            // 生成RSA密钥对
            KeyPair keyPair = generateRSAKeyPair();

            // 获取公钥和私钥
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();

            // 明文
            String plainText = "Hello, RSA!";

            // 加密
            byte[] encryptedBytes = encryptRSA(plainText, (RSAPublicKey) publicKey);

            // 解密
            byte[] bytes = Base64.getDecoder().decode(new String(encryptedBytes));
            String decryptedText = decryptRSA(bytes, (RSAPrivateKey) privateKey);

            // 输出结果
            System.out.println("原始文本：" + plainText);
            System.out.println("加密后：" + new String(encryptedBytes));
            System.out.println("解密后：" + decryptedText);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // 生成RSA密钥对
    public static KeyPair generateRSAKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(RSA);
        keyPairGenerator.initialize(2048);
        return keyPairGenerator.generateKeyPair();
    }

    // 使用公钥加密
    public static byte[] encryptRSA(String plainText, RSAPublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance(RSA, new BouncyCastleProvider());
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return rsaSplitCode(cipher, Cipher.DECRYPT_MODE, plainText.getBytes(),
                publicKey.getModulus().bitLength());
    }

    // 使用私钥解密
    public static String decryptRSA(byte[] encryptedBytes, RSAPrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance(RSA, new BouncyCastleProvider());
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return new String(rsaSplitCode(cipher, Cipher.DECRYPT_MODE, encryptedBytes,
                privateKey.getModulus().bitLength()), StandardCharsets.UTF_8);
    }

    /**
     * RSA分段加解密
     *
     * @param cipher  密码
     * @param mode    模式
     * @param data    数据
     * @param keySize Key大小
     * @return {@link byte[]}
     */
    static byte[] rsaSplitCode(Cipher cipher, int mode, byte[] data, int keySize) {
        // 最大块
        int maxBlock = 0;
        switch (mode) {
            case Cipher.ENCRYPT_MODE:
                maxBlock = keySize / Byte.SIZE - 66;
                break;
            case Cipher.DECRYPT_MODE:
                maxBlock = keySize / Byte.SIZE;
                break;
            default:
                throw new IllegalStateException(WebUtil.getMessage(MsgCd.ERR_MSG_BSS_2113870405) + mode);
        }
        ByteArrayOutputStream ctStream = null;
        try {
            ctStream = new ByteArrayOutputStream();
            int off = 0;
            try {
                while (off < data.length) {
                    int toCrypt = Math.min(maxBlock, data.length - off);
                    byte[] partialCT = cipher.doFinal(data, off, toCrypt);
                    ctStream.write(partialCT);
                    off += toCrypt;
                }
            } catch (Exception e) {
                log.error("Encryption and decryption threshold:{}, error info: ", off, e);
            }
            byte[] result = ctStream.toByteArray();
            IoUtil.close(ctStream);
            return result;
        } catch (Exception e) {
            log.error(e.getMessage());
        } finally {
            IoUtil.close(ctStream);
        }
        return null;
    }

    // 分段加密方法
    private static byte[] segmentEncrypt(Cipher cipher, byte[] plaintext) throws Exception {
        int blockSize = cipher.getBlockSize();
        byte[] encryptedData = new byte[0];
        if (blockSize == 0) {
            return new byte[0];
        }

        // 分段加密
        for (int i = 0; i < plaintext.length; i += blockSize) {
            int length = Math.min(blockSize, plaintext.length - i);
            byte[] block = new byte[length];
            System.arraycopy(plaintext, i, block, 0, length);
            byte[] encryptedBlock = cipher.update(block);
            encryptedData = concatenateArrays(encryptedData, encryptedBlock);
        }

        // 处理最后一个块
        byte[] finalBlock = cipher.doFinal();
        encryptedData = concatenateArrays(encryptedData, finalBlock);

        return encryptedData;
    }

    // 将两个字节数组连接起来
    private static byte[] concatenateArrays(byte[] array1, byte[] array2) {
        byte[] result = new byte[array1.length + array2.length];
        System.arraycopy(array1, 0, result, 0, array1.length);
        System.arraycopy(array2, 0, result, array1.length, array2.length);
        return result;
    }
}
